/*
 * Copyright (c) 2013-2018 上海汇数数据. All rights reserved.
 * @(#) SymmetricCipherApi.java 2018-08-06 17:26
 */

package cn.jh.common.core.crypto;

import org.jetbrains.annotations.Nullable;

import javax.crypto.Cipher;
import javax.crypto.spec.IvParameterSpec;
import java.security.Key;
import java.security.SecureRandom;
import java.security.spec.AlgorithmParameterSpec;

import static cn.jh.common.core.crypto.CipherKt.BITS_PER_BYTE;


/**
 * 基于JCA的对称密钥算法的 {@code CipherApi} 接口的基础实现
 *
 * @author Fuchun
 * @since 1.0
 */
public abstract class SymmetricCipherApi extends JcaCipherApi {

    /** 初始化向量的大小。*/
    int initializationVectorSize = 0;

    public SymmetricCipherApi(SymmetricCipherParams params) {
        super(params);
    }

    @Override
    public SymmetricCipherParams getParams() {
        return (SymmetricCipherParams) super.getParams();
    }

    protected String getTransformationString() {
        return getAlgorithm();
    }

    protected boolean isGenerateIV(boolean streaming) {
        return getParams().isGenerateIv();
    }

    /**
     * 生成 IV。
     *
     * @param streaming 是否基于流。
     * @return 新生成的 IV。
     */
    protected byte[] generateIV(boolean streaming) {
        int sizeInBytes = ensureIvSize();
        byte[] ivBytes = new byte[sizeInBytes];
        SecureRandom random = ensureSecureRandom();
        random.nextBytes(ivBytes);
        return ivBytes;
    }

    protected int ensureIvSize() {
        int size = getParams().getIvSize();
        if (size <= 0) {
            String msg = "initializationVectorSize property must be greater than zero. " +
                    "This number is typically set in the " + CipherApi.class.getSimpleName() +
                    " subclass constructor.  Also check your configuration to ensure that " +
                    "if you are setting a value, it is positive.";
            throw new IllegalStateException(msg);
        } else if (size % BITS_PER_BYTE != 0) {
            throw new IllegalStateException(
                    "initializationVectorSize property must be a multiple of 8 to represent as a byte array.");
        }
        return size / BITS_PER_BYTE;
    }

    protected SecureRandom ensureSecureRandom() {
        SecureRandom sr = getSecureRandom();
        if (sr == null) {
            sr = CipherKt.getDefaultSecurRandom();
        }
        return sr;
    }

    /**
     * 根据密钥获取固定的 IV。
     *
     * @param key 密钥。
     * @return 根据密钥获得的 IV。
     */
    protected byte[] getFixedIvByKey(byte[] key) {
        int ivSize = ensureIvSize();
        byte[] ivBytes = new byte[ivSize];
        System.arraycopy(key, 0, ivBytes, 0, ivSize);
        return ivBytes;
    }

    @Override
    public byte[] encrypt(byte[] plaintext, byte[] key) throws CryptoException {
        byte[] ivBytes = resolveIv(key);
        return encrypt(plaintext, key, ivBytes);
    }

    @Override
    public byte[] encrypt(byte[] plaintext, Key key) throws CryptoException {
        byte[] ivBytes = resolveIv(key.getEncoded());
        return encrypt(plaintext, key, ivBytes);
    }

    private byte[] resolveIv(byte[] key) {
        byte[] ivBytes = null;
        if (getParams().isUseFixedIvByKey()) {
            ivBytes = getFixedIvByKey(key);
        } else if (isGenerateIV(false)) {
            ivBytes = generateIV(false);
            if (ivBytes.length == 0) {
                throw new IllegalStateException("Initialization vector generation is enabled - " +
                        "generated vector cannot be null or empty.");
            }
        }
        return ivBytes;
    }

    @Override
    protected byte[] encrypt(byte[] plaintext, byte[] key, @Nullable byte[] iv) {
        int mode = Cipher.ENCRYPT_MODE;
        SymmetricCipherParams params = getParams();
        byte[] output;
        if (params.isPrependIv() && iv != null && iv.length > 0) {
            byte[] encrypted = cryptWithKeyBytes(plaintext, key, iv, mode);

            output = copyCipherText(encrypted, iv);
        } else {
            output = cryptWithKeyBytes(plaintext, key, iv, mode);
        }
        if (logger.isTraceEnabled()) {
            logger.trace("Incoming plaintext of size {}. Cipher-text byte array is size {}.",
                    plaintext.length, output.length);
        }
        return output;
    }

    @Override
    protected byte[] encrypt(byte[] plaintext, Key key, @Nullable byte[] iv) {
        int mode = Cipher.ENCRYPT_MODE;
        SymmetricCipherParams params = getParams();
        byte[] output;
        if (params.isPrependIv() && iv != null && iv.length > 0) {
            byte[] cipherText = cryptWithKey(plaintext, key, iv, mode);

            output = copyCipherText(cipherText, iv);
        } else {
            output = cryptWithKey(plaintext, key, iv, mode);
        }
        if (logger.isDebugEnabled()) {
            logger.debug("Incoming plaintext of size {}. Cipher-text byte array is size {}.",
                    plaintext.length, output.length);
        }
        return output;
    }

    @Override
    protected byte[] decrypt(byte[] cipherText, byte[] key, @Nullable byte[] parameters) {
        return super.decrypt(cipherText, key, parameters);
    }

    @Override
    public byte[] decrypt(byte[] cipherText, byte[] key) {
        SymmetricCipherParams params = getParams();

        Pair<byte[], byte[]> pair = resolveTextAndIv(params, cipherText, key);

        return decrypt(pair.second, key, pair.first);
    }

    @Override
    public byte[] decrypt(byte[] cipherText, Key key) throws CryptoException {
        SymmetricCipherParams params = getParams();

        Pair<byte[], byte[]> pair = resolveTextAndIv(params, cipherText, key.getEncoded());

        return decrypt(pair.second, key, pair.first);
    }

    private Pair<byte[], byte[]> resolveTextAndIv(SymmetricCipherParams params, byte[] cipherText, byte[] key) {
        byte[] iv = null, encrypted = cipherText;
        if (params.isUseFixedIvByKey()) {
            iv = getFixedIvByKey(key);
        } else if (params.isPrependIv() && isGenerateIV(false)) {
            try {
                int ivSize = params.getIvSize();
                int ivByteSize = ivSize / BITS_PER_BYTE;
                //now we know how large the iv is, so extract the iv bytes:
                iv = new byte[ivByteSize];
                System.arraycopy(cipherText, 0, iv, 0, ivByteSize);

                //remaining data is the actual encrypted cipherText.  Isolate it:
                int encryptedSize = cipherText.length - ivByteSize;
                encrypted = new byte[encryptedSize];
                System.arraycopy(cipherText, ivByteSize, encrypted, 0, encryptedSize);
            } catch (Exception ex) {
                throw new CryptoException("Unable to correctly extract the Initialization Vector or cipherText.");
            }
        }
        return new Pair<>(iv, encrypted);
    }


    @Override
    protected AlgorithmParameterSpec toParameterSpec(byte[] parameters) {
        return new IvParameterSpec(parameters);
    }

    private byte[] copyCipherText(byte[] cipherText, byte[] iv) {

        byte[] output = new byte[iv.length + cipherText.length];

        // copy iv to output
        System.arraycopy(iv, 0, output, 0, iv.length);
        // copy encrypted to output
        System.arraycopy(cipherText, 0, output, iv.length, cipherText.length);
        return output;
    }

    public int getInitializationVectorSize() {
        return initializationVectorSize;
    }

    public void setInitializationVectorSize(int initializationVectorSize) {
        if (initializationVectorSize % BITS_PER_BYTE != 0) {
            throw new IllegalArgumentException("Initialization vector sizes are specified in bits, " +
                    "but must be a multiple of 8 so they can be easily represented as a byte array.");
        }
        this.initializationVectorSize = initializationVectorSize;
    }

    private static class Pair<F, S> {
        private final F first;
        private final S second;

        Pair(F first, S second) {
            this.first = first;
            this.second = second;
        }
    }
}
